package com.avalieme.treinamento;

import java.io.BufferedReader;
import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStreamReader;

import weka.classifiers.bayes.NaiveBayesUpdateable;
import weka.core.SerializationHelper;

import com.avalieme.classificadores.ClassificadorException;
import com.avalieme.classificadores.Classificador;
import com.avalieme.util.Constantes;

public class TreinoAvalieMe {

	public static void main(String[] args) throws IOException, ClassificadorException {
		
		if (args.length != 2) {
			System.err.println("Erro. Os parametros corretos sao: "
					+ "<dirArquivos> <classe> "
					+ "Onde classe=feliz|triste|neutro");
			return;
		}

		String diretorioArquivos = args[0];
		String classe = args[1];

		Classificador classificador;
		try {
			if (classe.equals(Constantes.CLASSE_OPINATIVO) || classe.equals(Constantes.CLASSE_NEUTRO)) {
				classificador = (Classificador) SerializationHelper.read(Constantes.getModelo());
			} else {
				classificador = (Classificador) SerializationHelper.read(Constantes.getModeloPN());
			}
		} catch (Exception e) {
			if (classe.equals(Constantes.CLASSE_OPINATIVO) || classe.equals(Constantes.CLASSE_NEUTRO)) {
				classificador = new Classificador(new NaiveBayesUpdateable(),Constantes.getModelo(), false);
			} else {
				classificador = new Classificador(new NaiveBayesUpdateable(),Constantes.getModeloPN(), true);
			}
		}

		File diretorio = new File(diretorioArquivos);
		
		System.out.println("Comecar a treinar o classificador do tipo <"
				+ "Naive Bayes " + "> com os arquivos do diretorio <"
				+ diretorioArquivos + ">. Os tweets sao do tipo <" + classe
				+ ">.");

		for (File f : diretorio.listFiles()) {
			if (f.isHidden() || !f.getName().endsWith(".in")) {
				continue;
			}
			System.out.println("Treinando o arquivo atual " + f.getName());
			BufferedReader bf = new BufferedReader(new InputStreamReader(
					new FileInputStream(f), "UTF-8"));
			while (true) {
				String tweet = bf.readLine();
				if (tweet == null) {
					break;
				}
				tweet = tweet.toLowerCase();
				classificador.treinar(tweet, classe);
			}
			System.out.println("Fim do treino com o arquivo " + f.getName());
		}

		try {
			System.out.println("Salvando o modelo treinado ("
					+ classificador.getLocalModelo() + ")");
			SerializationHelper.write(classificador.getLocalModelo(),
					classificador);
			System.out.println("Modelo salvo com sucesso!!");
		} catch (Exception e) {
			e.printStackTrace();
		}
	}

}
